'''
Hand-eye calibration for Vive Tracker.

将vive固定于夹爪末端, 然后启动脚本. 这个脚本会启动 UR5 teach mode, 
将末端移动到随机采样点，点击 capture 按钮采集数据. 当收集足够的数据后，
点击 calibrate 按钮进行标定. 点击 save 按钮保存结果.

Author: Jin Liu.
Date: 2025/06/17
'''
from pathlib import Path

from ezrobot.robot import UR5
from ezrobot.sensor import ViveTracker, SensorManager
from ezrobot.sim import RobotClient
from ezrobot.utils import HandEyeConfig, solve_handeye
from ezgl import tb

BASE_DIR = Path(__file__).parent


if __name__ == '__main__':
    robot = UR5(arm_ip="192.168.77.36", with_gripper=False)
    robot.teach_mode()

    tracker = SensorManager(ViveTracker(device_index=1), frequency=30, buffer_size=10, streaming=True)
    client = RobotClient()

    handeye_config = HandEyeConfig()
    ee_poses = []
    vive_poses = []

    def capture():
        ee_poses.append(robot.get_pose())
        vive_poses.append(tracker.get_data())
        print(f"Captured {len(ee_poses)} samples")

        # visualize
        tb.set_value("samples cnt", value=f"samples cnt: {len(ee_poses)}")
        client.add_pose_axis(ee_poses[-1])
        client.add_pose_axis(handeye_config.base2map * vive_poses[-1] * handeye_config.hand2eye.inverse())

    def calibrate():
        hand2eye, residual = solve_handeye(ee_poses, [p.inverse() for p in vive_poses])
        base2map, residual1 = solve_handeye([p.inverse() for p in ee_poses], vive_poses)
        print(f"Calibration residual: {residual}, {residual1}")
        handeye_config.hand2eye = hand2eye
        handeye_config.base2map = base2map
        
        # visualize
        client.clear_items()
        for i in range(len(ee_poses)):
            client.add_pose_axis(handeye_config.base2map * vive_poses[i] * handeye_config.hand2eye.inverse())
            client.add_pose_axis(ee_poses[i])

    def on_timeout():
        # visualize
        joints = robot.get_joints_rad()
        client.set_joint(joints)

    def save_results():
        handeye_config.save(BASE_DIR / "../configs/handeye_config.yaml")

    with tb.window("control", None, 10, size=(500, 500), pos=(400, 0), frameless=False):
        tb.add_text("samples cnt", text="samples cnt: 0")
        tb.add_button("capture", callback=capture)
        tb.add_key_binding(tb.Key.Key_Space, callback=capture)
        tb.add_button("calibrate", callback=calibrate)
        tb.add_button("save", callback=save_results)
        tb.add_stretch(1)

        tb.add_timer("update", interval_ms=20, callback=on_timeout)

    tb.exec()


